import re
import base64
import binascii
import base91  # 需要安装: pip install pybase91
import base58  # 需要安装: pip install base58
from abstract_decoder import AbstractDecoder

class BaseDecoder(AbstractDecoder):

    def __init__(self):
        super().__init__("base")

    """
    多功能Base编码解码器
    支持: Base16, Base32, Base58, Base64, Base85, Base91
    """
    # 定义各种Base编码的特征和字符集
    BASE16_CHARS = "0123456789ABCDEFabcdef"
    BASE32_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567="
    BASE58_CHARS = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
    BASE64_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
    BASE64_URL_CHARS = (
        "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
    )
    BASE85_CHARS = "".join(chr(x) for x in range(33, 118))
    BASE91_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!#$%&()*+,-./:;<=>?@[]^_`{|}~"

    def detect_encoding(self, data) -> str:
        """
        检测输入的字符串使用了哪种Base编码

        返回: 'base16', 'base32', 'base58', 'base64', 'base85', 'base91' 或 'unknown'
        """
        # 转换为字符串并去除可能的空白字符
        s = str(data).strip()

        # 检查Base16
        if re.fullmatch(r"^[0-9A-Fa-f]+$", s) and len(s) % 2 == 0:
            return "base16"

        # 检查Base32
        s_upper = s.upper()
        if all(c in self.BASE32_CHARS for c in s_upper):
            # 检查填充是否符合Base32规范
            pad_count = s_upper.count("=")
            if pad_count == 0 or (len(s_upper) % 8 == 0 and pad_count <= 6):
                return "base32"

        # 检查Base58
        if all(c in self.BASE58_CHARS for c in s):
            return "base58"

        # 检查Base64
        s_clean = s.replace("=", "")
        if all(c in self.BASE64_CHARS for c in s_clean) and (
            len(s) % 4 == 0 or not s.endswith("=")
        ):
            return "base64"
        # 检查URL安全的Base64
        if all(c in self.BASE64_URL_CHARS for c in s):
            return "base64"

        # 检查Base85
        if all(33 <= ord(c) <= 117 for c in s):
            # 检查Adobe风格的包装
            if (s.startswith("<~") and s.endswith("~>")) or "z" in s:
                return "base85"
            # 检查标准Base85
            if len(s) % 5 == 0:
                return "base85"

        # 检查Base91
        if all(c in self.BASE91_CHARS for c in s):
            return "base91"

        return "unknown"

    def is_valid(self, data, key = None) -> bool:
        """
        检测输入的字符串是否是有效的Base编码
        参数:
        - data: 待检测的字符串
        - key: 解码密钥，用于某些编码需要提供
        返回: True 或 False
        """
        s = data.strip()
        # 自动检测编码类型
        encoding = self.detect_encoding(s)
        if encoding == "unknown":
            return False
        return True

    def decode(self, data, key = None) -> bytes:
        """
        解码输入的字符串
        参数:
        - data: 待解码的字符串
        - key: 解码密钥，用于某些编码需要提供
        返回: 解码后的字节串
        """
        s = data.strip()
        # 自动检测编码类型
        encoding = self.detect_encoding(s)
        if encoding == "unknown":
            raise ValueError("Invalid encoding")
        # 调用相应的解码函数
        if encoding == "base16":
            return self._decode_base16(s)
        elif encoding == "base32":
            return self._decode_base32(s)
        elif encoding == "base58":
            return self._decode_base58(s)
        elif encoding == "base64":
            result = self._decode_base64(s)
            return result
        elif encoding == "base85":
            return self._decode_base85(s)
        elif encoding == "base91":
            return self._decode_base91(s)
        else:
            raise ValueError("Invalid encoding")

    def _decode_base16(self, s: str) -> bytes:
        """Base16解码"""
        return binascii.unhexlify(s)

    def _decode_base32(self, s: str) -> bytes:
        """Base32解码"""
        # 添加必要的填充
        s += "=" * ((8 - len(s) % 8) % 8)
        result =  base64.b32decode(s)
        return result

    def _decode_base58(self, s: str) -> bytes:
        """Base58解码"""
        return base58.b58decode(s)

    def _decode_base64(self, s: str) -> bytes:
        """Base64解码"""
        s = s.strip()
        
        # 尝试标准Base64解码
        try:
            # 添加可能缺少的填充
            s_std = s + '=' * ((4 - len(s) % 4) % 4)
            return base64.b64decode(s_std)
        except (binascii.Error, ValueError):
            pass
        
        # 尝试URL安全Base64解码
        try:
            # 替换URL安全字符并添加填充
            # padding = len(s) % 4
            # if padding:
            #     s += '=' * (4 - padding)
            return base64.urlsafe_b64decode(s)
        except (binascii.Error, ValueError):
            pass
        
        # 尝试无填充解码
        try:
            # 先尝试标准无填充
            return base64.b64decode(s, validate=False)
        except (binascii.Error, ValueError):
            try:
                # 再尝试URL安全无填充
                s_url = s.replace('-', '+').replace('_', '/')
                return base64.b64decode(s_url, validate=False)
            except (binascii.Error, ValueError) as e:
                raise ValueError("无法解码为任何已知的Base64变体") from e

    def _decode_base85(self, s: str) -> str:
        """Base85解码"""
        try:
        # 处理Adobe风格的包装
            if s.startswith('<~') and s.endswith('~>'):
                s = s[2:-2]
            return base64.a85decode(s, adobe=True)
        except (binascii.Error, ValueError):
            try:
                # 尝试不带Adobe包装的版本
                return base64.a85decode(s, adobe=False)
            except (binascii.Error, ValueError):
                raise ValueError("Invalid Base85 string")

    def _decode_base91(self, s: str) -> str:
        """Base91解码"""
        return base91.decode(s)

if __name__ == "__main__":
    # 测试
    decoder = BaseDecoder()
    codes = [
        ("48656C6C6F20776F726C64", "base16"),
        ("JBSWY3DPEB3W64TMMQ", "base32"),
        ("JxF12TrwXzT5jvT", "base58"),
        ("SGVsbG8gd29ybGQ=", "base64"),
        # ("SGVsbG8gV29ybGQ-", "base64_url"),
        # ("87cURD]j7BEbo7", "base85"),
        (">OwJh>Io2Tv!lE", "base91"),
    ]
    for code, encoding in codes:
        result = decoder.decode(code)
        print(f"Decode {code} as {encoding}: {result}")
        assert result == b"Hello world"
    
    try:
        decoder.decode("Hello world")
    except ValueError as e:
        assert str(e) == "Invalid encoding" 
